
from genBase import genBase
from primitiveBase import * 
from hashlib import md5
import string
import tempfile 
import os

class genSizedString(genBase):
    '''
    scan input for strings and look for sizes of those strings
    create fuzzer that ties the strings to those detected sizers
    '''

    def __init__(self):
        '''
        Constructor
        '''

    def writeFuzzScripts(self, scriptFormat, outputDir, inputBuff=None):
        if inputBuff == None and self.scriptLegos == []:
            return
        if inputBuff != None:
            self.scriptLegos = []
            self.analyzeInput(inputBuff)
        if self.scriptLegos != []:
            print "creating scripts..."
            self.createdScripts = []
            for (foundString, sizers) in self.scriptLegos:
                if len(sizers) == 0:
                    fuzzedBuffOffsets = set(range(foundString.getStart(), foundString.getEnd()))
                    staticLego = static()
                    staticLegos = staticLego.buildAllStatics(fuzzedBuffOffsets, inputBuff) #TODO improve performance by only calling this minimal times
                    legos = []
                    legos.extend(staticLegos)
                    legos.append(foundString)
                    legos.sort()
                    scriptCode = "".join(map(str,legos))
                    scriptCode = scriptFormat % (scriptCode, self.fileId)
                    self.createdScripts.append(self.writeToFile(scriptCode, outputDir))
                else:
                    for foundSize in sizers:
                        fuzzedBuffOffsets = set()
                        sizerLocationSet = set( range(foundSize.getStart(), foundSize.getEnd()) )
                        stringLocationSet = set( range(foundString.getStart(), foundString.getEnd()) )
                        fuzzedBuffOffsets.update( sizerLocationSet )
                        fuzzedBuffOffsets.update( stringLocationSet )
                        staticLego = static()
                        staticLegos = staticLego.buildAllStatics(fuzzedBuffOffsets, inputBuff) #TODO improve performance by only calling this minimal times
                        legos = []
                        legos.extend(staticLegos)
                        legos.append(foundSize)
                        legos.append(foundString)
                        tempBlock = block()
                        createdBlocks = tempBlock.createBlocks(foundString)
                        legos.extend(createdBlocks)
                        legos.sort()
                        scriptCode = "".join(map(str,legos))
                        scriptCode = scriptFormat % (scriptCode, self.fileId)
                        self.createdScripts.append(self.writeToFile(scriptCode, outputDir))
 
                    #run consistency check that all regions are covered with a lego
        return self.createdScripts
    
    def analyzeInput(self, inputBuff):    
        foundStrings = self.findStringsInBinary(inputBuff)
        i = 0
        for foundString in foundStrings:
            print "finding sizers: %d / %d" % (i, len(foundStrings))
            i += 1
            sizers = self.findSizers(inputBuff, foundString.getStart(), foundString.getEnd(), foundString.getBlockName())
            self.scriptLegos.append((foundString, sizers))
        return self.scriptLegos
    
    def findStringsInBinary(self, inputBuff):
        if len(inputBuff) < self.MINSTRINGSIZE:
            return []
        stringsFound = []
        stringsFoundLoc = []

        for i in range(0, len(inputBuff)-self.MINSTRINGSIZE):
            #for this loop we dont want to re-analyze sections after they are already in the stringsFound list
            #other search loops may want to re-analyze sections already in the results list
            skipIndex = False
            for (start, end) in stringsFoundLoc:
                if start <= i and end >= i:
                    skipIndex = True
                    break
            if skipIndex:
                continue
            
            if inputBuff[i] in string.ascii_letters and inputBuff[i+self.MINSTRINGSIZE] in string.ascii_letters:
                stringFound = True
                for x in range(i+1, i+self.MINSTRINGSIZE-1):
                    if inputBuff[x] not in string.ascii_letters:
                        stringFound = False
                if stringFound:
                    stringStart = i
                    stringEnd = -1
                    #find the end of the string, possibly adding substrings to the list based on delims other than non-ascii chars
                    for x in range(i, len(inputBuff)):
                        if inputBuff[x] not in string.ascii_letters:
                            stringEnd = x
                            break
                    if x == len(inputBuff)-1:
                        stringEnd = x+1
                    if stringEnd == -1:
                        raise Exception('stringEnd not found in findStringsInBinary')
                    m = md5()
                    m.update(inputBuff[stringStart:stringEnd])
                    name = m.hexdigest()
                    foundString = stringPrimitive(str(inputBuff[stringStart:stringEnd]), stringStart, name=name)
                    stringsFound.append(foundString)
                    stringsFoundLoc.append((stringStart, stringEnd))
                    print "Found String: %s" % (str(inputBuff[stringStart:stringEnd]))
                    
        return stringsFound
            
